﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.Utils;
using Sets = Algoritmia.DataStructures.Sets;
using Algoritmia.DataStructures.Queues;

namespace Algoritmia.StateSpaces
{
    public class StateSpaceTopsorter<TState,TDecision>
    {
        private readonly Func<Sets.ISet<TState>> _createSet;
        private readonly Func<ILifo<Tuple<TState, IEnumerable<TDecision>>>> _createLifo;

        public StateSpaceTopsorter(Func<ILifo<Tuple<TState, IEnumerable<TDecision>>>> createLifo = null,
                                   Func<Sets.ISet<TState>> createSet = null )
        {
            _createLifo = createLifo ?? (() => AlgoritmiaFactory.Lifos<Tuple<TState, IEnumerable<TDecision>>>.Default);
            _createSet = createSet ?? (() => AlgoritmiaFactory.Sets<TState>.Default);
        }

        #region Private PostorderTraverses

        private IEnumerable<TState> PostorderTraverse(TState s, IBackwardsStateSpace<TState,TDecision> space, Sets.ISet<TState> visited,
            ILifo<Tuple<TState,IEnumerable<TDecision>>> stack)
        {
            visited.Add(s);
            stack.Push(new Tuple<TState, IEnumerable<TDecision>>(s,
                                                                 space.IncomingDecisions(s)));
            while (stack.Count > 0)
            {
                var tp = stack.Pop();
                var u = tp.Item1;
                var decisions = tp.Item2;
                var q = decisions.Where(d => !visited.Contains(space.Undo(u, d))).ToList();
                if (q.Count > 0)
                {
                    var w = space.Undo(u, q[0]);
                    stack.Push(new Tuple<TState, IEnumerable<TDecision>>(u, q));
                    stack.Push(new Tuple<TState, IEnumerable<TDecision>>(w, space.IncomingDecisions(w)));
                    visited.Add(w);
                }
                else
                {
                    yield return u;
                }
            }
        }

        private IEnumerable<TState> PostorderTraverse(TState s, IForwardStateSpace<TState, TDecision> space, Sets.ISet<TState> visited,
           ILifo<Tuple<TState, IEnumerable<TDecision>>> stack)
        {
            visited.Add(s);
            stack.Push(new Tuple<TState, IEnumerable<TDecision>>(s,
                                                                 space.Decisions(s)));
            while (stack.Count > 0)
            {
                var tp = stack.Pop();
                var u = tp.Item1;
                var decisions = tp.Item2;
                var q = decisions.Where(d => !visited.Contains(space.Decide(u, d))).ToList();
                if (q.Count > 0)
                {
                    var w = space.Decide(u, q[0]);
                    stack.Push(new Tuple<TState, IEnumerable<TDecision>>(u, q));
                    stack.Push(new Tuple<TState, IEnumerable<TDecision>>(w, space.Decisions(w)));
                    visited.Add(w);
                }
                else
                {
                    yield return u;
                }
            }
        }

        #endregion

        public IEnumerable<TState> Topsorted(IBackwardsStateSpace<TState,TDecision> space)
        {
            var visited = _createSet();
            var stack = _createLifo();

            foreach(var state in space.FinalStates)
            {
                if ( !visited.Contains(state))
                {
                    foreach (var s in PostorderTraverse(state, space, visited, stack))
                    {
                        yield return s;
                    }
                }
            }
        }

        public IEnumerable<TState> Topsorted(IForwardStateSpace<TState,TDecision> space)
        {
            var visited = _createSet();
            var stack = _createLifo();

            foreach (var state in space.InitialStates)
            {
                if (!visited.Contains(state))
                {
                    foreach (var s in PostorderTraverse(state, space, visited, stack).Reverse())
                    {
                        yield return s;
                    }
                }
            }
        }
    }
}
